import cv2
import numpy as np

class VisionProcessor:
    def __init__(self):
        # 棋盘标准参数（未旋转时的角点坐标）
        self.ref_corners = np.array([[0, 0], [300, 0], [300, 300], [0, 300]], dtype=np.float32)
        
    def detect_board(self, frame):
        """检测棋盘区域并返回旋转角度和变换后的棋盘图像"""
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        edges = cv2.Canny(blurred, 50, 150)
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 寻找最大轮廓（棋盘外框）
        max_area = 0
        board_contour = None
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > max_area:
                max_area = area
                board_contour = cnt
        
        # 提取角点并计算旋转角度
        if board_contour is not None:
            epsilon = 0.02 * cv2.arcLength(board_contour, True)
            approx = cv2.approxPolyDP(board_contour, epsilon, True)
            if len(approx) == 4:
                current_corners = np.array(approx, dtype=np.float32).reshape(4, 2)
                # 计算旋转矩阵
                matrix = cv2.getPerspectiveTransform(current_corners, self.ref_corners)
                # 提取旋转角度（假设绕中心旋转）
                dx = current_corners[1][0] - current_corners[0][0]
                dy = current_corners[1][1] - current_corners[0][1]
                angle = np.degrees(np.arctan2(dy, dx))
                return angle, matrix
        return 0.0, None

    def detect_pieces(self, frame, matrix):
        """检测棋子位置并返回3x3棋盘状态（0:空, 1:黑棋, -1:白棋）"""
        if matrix is None:
            return None
        # 透视变换校正棋盘
        warped = cv2.warpPerspective(frame, matrix, (300, 300))
        state = np.zeros((3, 3), dtype=int)
        cell_size = 100
        for i in range(3):
            for j in range(3):
                x, y = j*cell_size + 50, i*cell_size + 50
                roi = warped[y-40:y+40, x-40:x+40]
                # 通过颜色判断棋子
                hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
                black_lower = np.array([0, 0, 0])
                black_upper = np.array([180, 255, 50])
                white_lower = np.array([0, 0, 200])
                white_upper = np.array([180, 30, 255])
                
                # 统计黑色和白色像素
                mask_black = cv2.inRange(hsv, black_lower, black_upper)
                mask_white = cv2.inRange(hsv, white_lower, white_upper)
                if np.sum(mask_white) > 1000:
                    state[i][j] = -1
                elif np.sum(mask_black) > 1000:
                    state[i][j] = 1
        return state